#pragma once

#include <functional>
#include <list>
#include <vector>
#include <shared_mutex>
#include <memory>
#include <mutex>

template<class Key, class Value, class Hash = std::hash<Key>>
class threadsafe_lookup_table
{
private:
    class bucket_type
    {
    private:
        typedef std::pair<Key, Value> bucket_value;
        typedef std::list<bucket_value> bucket_data;
        typedef typename std::list<bucket_value>::iterator bucket_iterator;

    private:
        bucket_data data;
        mutable std::shared_mutex mutex;

        // 根据 key 返回对应节点的迭代器, 注意这里不能用 const 修饰, 否则返回的迭代器也会加上 const
        bucket_iterator find_entry_for(const Key& key)
        {
            return std::find_if(data.begin(), data.end(), [&](const bucket_value& item){
                return item.first == key;
            });
        }

    public:
        Value value_for(const Key& key, const Value& default_value) 
        {
            std::shared_lock<std::shared_mutex> lock(mutex);    // 共享读
            bucket_iterator found_entry = find_entry_for(key);
            return (found_entry == data.end()) ? default_value : found_entry->second;
        }

        void add_or_update_mapping(const Key& key, const Value& value)
        {
            std::unique_lock<std::shared_mutex> lock(mutex);    // 独占写
            bucket_iterator found_entry = find_entry_for(key);
            if (found_entry == data.end()) {
                data.push_back(bucket_value(key, value));
            }   
            else {
                found_entry->second = value;
            }
        }

        void remove_mapping(const Key& key)
        {
            std::unique_lock<std::shared_mutex> lock(mutex);    // 独占删
            bucket_iterator found_entry = find_entry_for(key);
            if (found_entry != data.end()) {
                data.erase(found_entry);
            }
        }
    };

private:
    std::vector<std::unique_ptr<bucket_type>> buckets;
    Hash hasher;

    bucket_type& get_bucket(const Key& key) const
    {
        std::size_t bucket_index = hasher(key) % buckets.size();
        return *buckets[bucket_index];
    }

public:
    typedef Key key_type;
    typedef Value mapped_type;
    typedef Hash hash_type;

    threadsafe_lookup_table(unsigned int num_buckets = 19, const Hash& hasher_ = Hash()) 
    : buckets(num_buckets), hasher(hasher_)
    {
        for (unsigned int i = 0; i < num_buckets; i++) {
            buckets[i].reset(new bucket_type);
        }
    }
    threadsafe_lookup_table(const threadsafe_lookup_table&) = delete;
    threadsafe_lookup_table& operator=(const threadsafe_lookup_table&) = delete;

    Value value_for(const Key& key, const Value& default_value = Value()) const
    {
        return get_bucket(key).value_for(key, default_value);
    }

    void add_or_update_mapping(const Key& key, const Value& value)
    {
        get_bucket(key).add_or_update_mapping(key, value);
    }

    void remove_mapping(const Key& key)
    {
        get_bucket(key).remove_mapping(key);
    }
};